from .geometry import get_normal_field, get_average_normal
import numpy as np

def get_force_norm(F: np.ndarray) -> np.ndarray:
    '''
    Get the norm of the 3D Forces.
    Args:
        F (np.ndarray): 3D Forces of shape (N, 400, 3) or (N, 20,20,3).
    Returns:
        np.ndarray: The norm of the 3D Forces of shape (N, 400) or (N, 20,20).
    '''
    return np.linalg.norm(F, axis=-1)

def get_xy_force(F: np.ndarray) -> np.ndarray:
    '''
    Get the XY components of the 3D Forces.
    Args:
        F (np.ndarray): 3D Forces of shape (N, 400, 3) or (N, 20,20,3).
    Returns:
        np.ndarray: The XY components of the 3D Forces of shape (N, 400, 2) or (N, 20,20,2).
    '''
    return np.linalg.norm(F[...,0:2], axis=-1)

def get_z_force(F: np.ndarray) -> np.ndarray:
    '''
    Get the Z component of the 3D Forces.
    Args:
        F (np.ndarray): 3D Forces of shape (N, 400, 3) or (N, 20,20,3).
    Returns:
        np.ndarray: The Z component of the 3D Forces of shape (N, 400) or (N, 20,20).
    '''
    return -F[...,2]

def force_normal_decompose(P: np.ndarray, F: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    '''
    Decompose the 3D Forces into normal and tangential components.
    Args:
        P (np.ndarray): 3D Points of shape (N, 400, 3).
        F (np.ndarray): 3D Forces of shape (N, 400, 3).
    Returns:
        Tuple[np.ndarray, np.ndarray]: The normal and tangential components of the 3D Forces, each of shape (N, 400, 3).
    '''
    assert P.shape == F.shape and P.shape[-1] == 3, "P and F must have the same shape and last dimension must be 3."
    assert P.shape[1:] == (400,3) or P.shape[1:] == (20,20,3), "P and F must have shape (N, 400, 3) or (N, 20,20, 3)."
    P = P.reshape(P.shape[0], -1, 3)  # (N, 400, 3)
    F = F.reshape(F.shape[0], -1, 3)  # (N, 400, 3)
    normal = get_normal_field(P)  # (N, 400, 3)
    Fn = np.sum(F * normal, axis=-1, keepdims=True) * normal  # (N, 400, 3)
    Ft = F - Fn  # (N, 400, 3)
    return Fn, Ft

def get_fr_normal_decomposition(P: np.ndarray, F: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    '''
    Get the normal force decomposition of the 3D resultant Forces into normal and tangential components.
    Args:
        P (np.ndarray): 3D Points of shape (N, 400, 3).
        F (np.ndarray): 3D Forces of shape (N, 400, 3).
    Returns:
        Tuple[np.ndarray, np.ndarray]: The normal and tangential components of the 3D Resultant Forces, each of shape (N, 3).
    '''
    average_normal = get_average_normal(P,F)  # (N, 3)
    Fr = np.sum(F, axis=-2)  # (N, 3)
    Fnr = np.sum(Fr * average_normal,axis=-1,keepdims=True) * average_normal  # (N, 3)
    Ftr = Fr - Fnr  # (N, 3)
    return Fnr, Ftr